"""
股票数据处理方案
读取stock_ticker.h5中的所有股票代码，获取从IPO到现在的日级别数据，
并计算各种均线和斜率指标
"""

import h5py
import akshare as ak
import pandas as pd
import numpy as np
import time
from datetime import datetime
import os
import sys
import talib
ROOT = os.path.dirname(os.path.abspath("."))
sys.path.append(ROOT)
import os.path as osp

def read_ticker_from_h5(file_path=osp.join(ROOT,'data','stock_ticker.h5')):
    """
    从h5文件中读取所有股票代码
    
    参数:
    file_path : str
        h5文件路径
    返回:
    list: 包含上海和深圳交易所所有股票代码的列表
    """
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"文件 {file_path} 不存在")
    
    tickers = []
    with h5py.File(file_path, 'r') as h5_file:
        # 读取上海交易所股票代码
        if 'sh_tickers' in h5_file:
            sh_tickers = [ticker.decode('utf-8') for ticker in h5_file['sh_tickers'][:]]
            tickers.extend(sh_tickers)
            print(f"读取到上海交易所 {len(sh_tickers)} 只股票")
        
        # 读取深圳交易所股票代码
        if 'sz_tickers' in h5_file:
            sz_tickers = [ticker.decode('utf-8') for ticker in h5_file['sz_tickers'][:]]
            tickers.extend(sz_tickers)
            print(f"读取到深圳交易所 {len(sz_tickers)} 只股票")
    
    print(f"总共读取到 {len(tickers)} 只股票代码")
    return tickers


def calculate_moving_averages(df, price_col='close', volume_col='volume'):
    """
    计算各种周期的均线和斜率
    
    参数:
    df : pd.DataFrame
        股票数据DataFrame
    price_col : str
        价格列名，默认为'close'
    volume_col : str
        成交量列名，默认为'volume'
        
    返回:
    pd.DataFrame: 添加了均线和斜率指标的DataFrame
    """
    # 定义需要计算的均线周期
    ma_periods = [7, 13, 30, 60, 90, 180, 300]
    
    df_result = df.copy()
    
    # 计算价格均线
    for period in ma_periods:
        # 计算价格均线
        ma_price_col = f'ma_price_{period}'
        df_result[ma_price_col] = df_result[price_col].rolling(window=period, min_periods=1).mean()
        
        # 计算价格均线斜率
        slope_price_col = f'slope_price_{period}'
        if len(df_result) >= period:
            df_result[slope_price_col] = df_result[ma_price_col].diff(period) / period
        else:
            df_result[slope_price_col] = 0
    
    # 计算成交量均线
    for period in ma_periods:
        # 计算成交量均线
        ma_volume_col = f'ma_volume_{period}'
        df_result[ma_volume_col] = df_result[volume_col].rolling(window=period, min_periods=1).mean()
        
        # 计算成交量均线斜率
        slope_volume_col = f'slope_volume_{period}'
        if len(df_result) >= period:
            df_result[slope_volume_col] = df_result[ma_volume_col].diff(period) / period
        else:
            df_result[slope_volume_col] = 0
    df_result = df_result.fillna(0)
    return df_result


def get_stock_data_with_retry(ticker, max_retries=3):
    """
    获取单只股票的历史数据，带重试机制
    
    参数:
    ticker : str
        股票代码
    max_retries : int
        最大重试次数
        
    返回:
    pd.DataFrame: 股票历史数据
    """
    # 转换股票代码格式
    if '.' in ticker:
        code_part = ticker.split('.')[0]
    else:
        code_part = ticker
    
    for attempt in range(max_retries):
        try:
            # 获取从IPO至今的日线数据
            stock_data = ak.stock_zh_a_hist(
                symbol=code_part,
                period="daily",
                adjust="qfq"  # 前复权
            )
            
            if stock_data.empty:
                print(f"股票 {ticker} 数据为空")
                return None
                
            # 重命名列以匹配标准格式
            stock_data = stock_data.rename(columns={
                '日期': 'date',
                '股票代码':'tic',
                '开盘': 'open',
                '收盘': 'close',
                '最高': 'high',
                '最低': 'low',
                '成交量': 'volume',
                '成交额': 'amount',
                '振幅': 'amplitude',
                '涨跌幅': 'change_rate',
                '涨跌额': 'change_amount',
                '换手率': 'turnover_rate',

            })
            
            # 选择需要的列
            stock_data = stock_data[['date', 'open', 'high', 'low', 'close', 'volume', 'amount', 'amplitude', 'change_rate', 'change_amount', 'turnover_rate']]
            
            # 添加股票代码列
            stock_data['tic'] = ticker
            
            print(f"成功获取股票 {ticker} 的历史数据，共 {len(stock_data)} 条记录")
            return stock_data
            
        except Exception as e:
            print(f"第 {attempt + 1} 次尝试获取 {ticker} 数据失败: {e}")
            if attempt < max_retries - 1:
                time.sleep(2 ** attempt)  # 指数退避
            else:
                print(f"获取 {ticker} 数据最终失败")
                return None


def process_all_stocks(ticker_file=osp.join(ROOT,'data','stock_ticker.h5'), output_dir=osp.join(ROOT,'data','stock_daily_data')):
    """
    处理所有股票数据：读取代码、获取数据、计算指标并保存
    
    参数:
    ticker_file : str
        股票代码h5文件路径
    output_dir : str
        输出目录
    """
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"创建输出目录: {output_dir}")

    # 读取股票代码
    tickers = read_ticker_from_h5(ticker_file)
    
    # 处理每只股票
    processed_count = 0
    error_count = 0
    total_tickers = len(tickers)
    
    print(f"开始处理 {total_tickers} 只股票...")
    
    for i, ticker in enumerate(tickers):
        try:

            # 判断文件是否存在 如果存在直接进入下一个
            if os.path.exists(os.path.join(output_dir, f"{ticker.replace('.', '_')}.csv")):
                #打印 跳过
                print(f"文件已存在，跳过处理股票 {ticker}")
                continue

            # 打印时间
            timestr=time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
            print(f"{timestr} 处理进度: {i + 1}/{total_tickers} ({(i + 1) / total_tickers * 100:.2f}%) - 处理股票 {ticker}")
            # 获取股票数据
            stock_data = get_stock_data_with_retry(ticker)
            
            if stock_data is None:
                error_count += 1
                continue
            
            # 计算均线和斜率指标
            processed_data = calculate_moving_averages(stock_data)

            #保存到CSV文件
            output_file = os.path.join(output_dir, f"{ticker.replace('.', '_')}.csv")
            processed_data.to_csv(output_file, index=False, encoding='utf-8-sig')

            processed_count += 1
            
            # 添加延迟避免请求过于频繁
            time.sleep(3)
            
        except Exception as e:
            print(f"处理股票 {ticker} 时出错: {e}")
            error_count += 1
            continue
    
    print(f"\n处理完成!")
    print(f"成功处理: {processed_count} 只股票")
    print(f"处理失败: {error_count} 只股票")
    print(f"输出目录: {os.path.abspath(output_dir)}")


def process_sample_stocks(ticker_file='stock_ticker.h5', output_dir='stock_data_sample', sample_size=5):
    """
    处理部分股票数据作为示例（用于测试）
    
    参数:
    ticker_file : str
        股票代码h5文件路径
    output_dir : str
        输出目录
    sample_size : int
        采样数量
    """
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"创建输出目录: {output_dir}")
    

    print("已设置Akshare超时时间为30秒")
    
    # 读取股票代码
    tickers = read_ticker_from_h5(ticker_file)
    
    # 只处理前sample_size只股票作为示例
    sample_tickers = tickers[:sample_size]
    
    print(f"开始处理 {len(sample_tickers)} 只股票作为示例...")
    
    processed_count = 0
    error_count = 0
    
    for i, ticker in enumerate(sample_tickers):
        try:
            print(f"处理进度: {i+1}/{len(sample_tickers)} - 处理股票 {ticker}")
            
            # 获取股票数据
            stock_data = get_stock_data_with_retry(ticker)
            
            if stock_data is None:
                error_count += 1
                continue
            
            # 计算均线和斜率指标
            processed_data = calculate_moving_averages(stock_data)
            
            # 保存到CSV文件
            output_file = os.path.join(output_dir, f"{ticker.replace('.', '_')}.csv")
            processed_data.to_csv(output_file, index=False, encoding='utf-8-sig')
            
            processed_count += 1
            
            # 添加延迟避免请求过于频繁
            time.sleep(0.5)
            
        except Exception as e:
            print(f"处理股票 {ticker} 时出错: {e}")
            error_count += 1
            continue
    
    print(f"\n示例处理完成!")
    print(f"成功处理: {processed_count} 只股票")
    print(f"处理失败: {error_count} 只股票")
    print(f"输出目录: {os.path.abspath(output_dir)}")


def choice_method():
    print("股票数据处理方案")
    print("=" * 50)
    print("1. 处理所有股票数据")
    print("2. 处理部分股票数据作为示例")

    choice = input("请选择处理方式 (1/2): ").strip()

    if choice == '1':
        process_all_stocks()
    elif choice == '2':
        sample_size = input("请输入示例股票数量 (默认5): ").strip()
        sample_size = int(sample_size) if sample_size.isdigit() else 5
        process_sample_stocks(sample_size=sample_size)
    else:
        print("无效选择，运行示例处理...")
        process_sample_stocks()


def add_other_indicator():
    # 读取股票代码
    tickers = read_ticker_from_h5()
    # 定义输出目录
    output_dir = osp.join(ROOT, 'data', 'stock_daily_data')
    
    # 遍历每只股票
    for i,ticker in enumerate(tickers):
        print(f"处理进度: {i + 1}/{len(tickers)} - 处理股票 {ticker}")
        try:
            # 构建股票数据文件路径
            file_path = os.path.join(output_dir, f"{ticker.replace('.', '_')}.csv")
            
            # 检查文件是否存在
            if not os.path.exists(file_path):
                print(f"股票 {ticker} 的数据文件不存在，跳过处理")
                continue
            
            # 读取股票数据
            df = pd.read_csv(file_path)
            
            # 计算MACD指标
            df['macd'], df['macdsignal'], df['macdhist'] = talib.MACD(df['close'], fastperiod=12, slowperiod=26, signalperiod=9)
            df = df.drop(columns=['macdsignal', 'macdhist'])  # 只保留macd列
            
            # 计算布林带指标
            df['boll_ub'], df['boll_mb'], df['boll_lb'] = talib.BBANDS(df['close'], timeperiod=20, nbdevup=2, nbdevdn=2, matype=0)
            df = df.drop(columns=['boll_mb'])  # 只保留上轨和下轨
            
            # 计算RSI指标
            df['rsi_30'] = talib.RSI(df['close'], timeperiod=30)
            
            # 计算CCI指标
            df['cci_30'] = talib.CCI(df['high'], df['low'], df['close'], timeperiod=30)
            
            # 计算DX指标
            df['dx_30'] = talib.DX(df['high'], df['low'], df['close'], timeperiod=30)
            df = df.fillna(0)
            # 保存更新后的数据
            df.to_csv(file_path, index=False, encoding='utf-8-sig')
            print(f"成功为股票 {ticker} 添加新指标")
            
        except Exception as e:
            print(f"为股票 {ticker} 添加指标时出错: {e}")
def main():
    """
    主函数
    """
    #choice_method()
    # process_all_stocks()
    # add_other_indicator()
    stock_data = get_stock_data_with_retry("600036")

if __name__ == "__main__":
    main()